package com.atguigu.ds.impl;

import com.atguigu.ds.Tree;

import java.util.function.Consumer;

public class BinarySearchTree<E extends Comparable<E>> implements Tree<E> {

    private int size ;

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size==0;
    }

    private static class Node<E extends Comparable<E>>{
        private E item ;
        private Node<E> leftChild ;
        private Node<E> rightChild ;
        Node(E item , Node<E> leftChild , Node<E> rightChild){
            this.item = item ;
            this.leftChild = leftChild ;
            this.rightChild = rightChild ;
        }
    }

    private Node<E> root ;

    @Override
    public void add(E element) {
        if(root==null){
            root = new Node(element,null,null);
        }else{
            //Node<E> pre = root ;
            Node<E> curr = root ;
            Node<E> pre = null ;
            while(curr!=null){
                if(element.compareTo(curr.item)<0){
                    pre = curr ;
                    curr = curr.leftChild;
                }else if(element.compareTo(curr.item)>0){
                    pre = curr ;
                    curr = curr.rightChild;
                }else{
                    return ;
                }
            }
            Node<E> node = new Node(element,null,null);
            if(element.compareTo(pre.item)<0){
                pre.leftChild = node  ;
            }else{
                pre.rightChild = node;
            }
        }
        size++;
    }

    @Override
    public E remove(E element) {

        //虚拟节点。如果待删除节点是根节点，则需要虚拟节点
        Node<E> pre = new Node<>(null,root,null);

        Node<E> curr = root ;
        //Node<E> pre = null ;
        while(curr!=null){
            if(element.compareTo(curr.item)<0){
                pre = curr ;
                curr = curr.leftChild ;
            }else if(element.compareTo(curr.item)>0){
                pre = curr ;
                curr = curr.rightChild;
            }else{
                break;
            }
        }

        //说明已经挖到叶子节点的子节点（叶子节点的子节点是null），说明没有找到待删除节点，删除方法直接结束即可
        if(curr==null){
            return null;
        }

        //如果代码能运行到这里，说明curr不等于null，也就是说明找到待删除节点了
        //下面分三种情况讨论：
        //① 待删除节点是叶子节点
        //② 待删除节点有一个子节点
        //③ 待删除节点有两个子节点

        //① 待删除节点是叶子节点
        if(curr.leftChild==null && curr.rightChild==null){
            if(pre.leftChild==curr){//待删除节点是它的父节点的左孩子
                pre.leftChild = null ;
            }else{
                pre.rightChild = null ;
            }
            size-- ;
        }else if(curr.leftChild==null || curr.rightChild==null){
            //② 待删除节点有一个子节点
            //currChildNode指的是待删除节点的子节点
            Node<E> currChildNode = curr.leftChild==null ? curr.rightChild : curr.leftChild;
            if(pre.leftChild==curr){//待删除节点是它的父节点的左孩子
                pre.leftChild = currChildNode ;
            }else{
                pre.rightChild = currChildNode ;
            }
            size-- ;
        }else{
            //③ 待删除节点有两个子节点
            //3-1 先找到右子树的最小节点，然后将其删除
            Node<E> minNode = curr.rightChild;
            while(minNode.leftChild!=null){
                minNode = minNode.leftChild ;
            }
            remove(minNode.item);
            //3-2 将minNode加载到待删除节点(curr节点)的位置
            minNode.leftChild = curr.leftChild;
            minNode.rightChild=curr.rightChild;
            //3-3 pre节点之前是指向curr的。现在需要指向minNode
            if(pre.leftChild==curr){
                pre.leftChild = minNode;
            }else{
                pre.rightChild = minNode ;
            }
        }

        //size--;

        //说明curr(待删除节点)是根节点。
        //通过上述代码，删除根节点，以及将右子树最小的节点替换到根节点这个位置，这些功能都已经完成
        //但是我们有一个成员变量-root。我们需要更新一下root的指向。否则add方法出错，否则迭代也会出错。因为root仍然指向之前被删除的那个根节点
        if(pre.item==null){
            root = pre.leftChild ;
        }

        return element;
    }

    @Override
    public void forEach(Consumer<E> consumer) {
        forEach(consumer,TRAVERSAL_ORDER.IN);
    }

    @Override
    public void forEach(Consumer<E> consumer, TRAVERSAL_ORDER order) {
        treeWalker(root,consumer,order);
    }

    private void treeWalker(Node<E> node , Consumer<E> consumer ,TRAVERSAL_ORDER order ){
        if(order==TRAVERSAL_ORDER.PRE){
            if(node!=null) {
                consumer.accept(node.item);
                treeWalker(node.leftChild, consumer, TRAVERSAL_ORDER.PRE);
                treeWalker(node.rightChild, consumer, TRAVERSAL_ORDER.PRE);
            }
        }else if(order==TRAVERSAL_ORDER.IN){
            if(node!=null) {
                treeWalker(node.leftChild, consumer, TRAVERSAL_ORDER.IN);
                consumer.accept(node.item);
                treeWalker(node.rightChild, consumer, TRAVERSAL_ORDER.IN);
            }
        }else{
            if(node!=null) {
                treeWalker(node.leftChild, consumer, TRAVERSAL_ORDER.POST);
                treeWalker(node.rightChild, consumer, TRAVERSAL_ORDER.POST);
                consumer.accept(node.item);
            }
        }
    }
}
